# Copyright (c) 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, List, Optional, overload

from nvidia.dali.types import DALIDataType

class TensorCPU:
    # @overload
    # def __init__(self, object: capsule, layout: str = ...) -> None: ...
    @overload
    def __init__(self, object, layout: str = ...) -> None: ...

    # @overload
    # def __init__(self, b: buffer, layout: str = ..., is_pinned: bool = ...) -> None: ...
    @overload
    def __init__(self, b, layout: str = ..., is_pinned: bool = ...) -> None: ...
    def _as_gpu(self, *args, **kwargs) -> Any: ...
    def as_cpu(self) -> TensorCPU: ...

    # def __dlpack__(self, stream: Optional[int] = None) -> capsule: ...

    def copy_to_external(self, ptr: object) -> None: ...
    def data_ptr(self) -> object: ...
    def get_property(self, arg0: str) -> object: ...
    def layout(self) -> str: ...
    def shape(self) -> list: ...
    def source_info(self) -> str: ...
    def squeeze(self, dim: object = ...) -> bool: ...
    @property
    def dtype(self) -> DALIDataType: ...
    @property
    def __array_interface__(self) -> dict: ...

class TensorGPU:
    # @overload
    # def __init__(self, object: capsule, layout: str = ...) -> None: ...
    @overload
    def __init__(self, object, layout: str = ...) -> None: ...

    # @overload
    # def __init__(self, object: object, layout: str = ..., device_id: int = ...) -> None: ...
    @overload
    def __init__(self, object, layout: str = ..., device_id: int = ...) -> None: ...

    # def __dlpack__(self, stream: Optional[int] = None) -> capsule: ...

    def as_cpu(self) -> TensorCPU: ...
    def copy_to_external(
        self,
        ptr: object,
        cuda_stream: object = ...,
        non_blocking: bool = ...,
        use_copy_kernel: bool = ...,
    ) -> None: ...
    def data_ptr(self) -> object: ...
    def get_property(self, arg0: str) -> object: ...
    def layout(self) -> str: ...
    def shape(self) -> list: ...
    def source_info(self) -> str: ...
    def squeeze(self, dim: object = ...) -> bool: ...
    @property
    def dtype(self) -> DALIDataType: ...
    @property
    def __cuda_array_interface__(self) -> dict: ...

class TensorListCPU:
    # __hash__: ClassVar[None] = ...

    # @overload
    # def __init__(self, object: capsule, layout: str = ...) -> None: ...
    @overload
    def __init__(self, object, layout: str = ...) -> None: ...
    @overload
    def __init__(self, tl: TensorListCPU, layout: object = ...) -> None: ...

    # @overload
    # def __init__(self, b: buffer, layout: str = ..., is_pinned: bool = ...) -> None: ...
    @overload
    def __init__(self, b, layout: str = ..., is_pinned: bool = ...) -> None: ...
    @overload
    def __init__(self, list_of_tensors: list, layout: str = ...) -> None: ...

    # def _as_gpu(self, *args, **kwargs) -> TensorListGPU:  ...
    def as_cpu(self) -> TensorListCPU: ...

    # def as_array(self) -> numpy.ndarray: ...
    def as_array(self) -> Any: ...
    def as_reshaped_tensor(self, arg0: List[int]) -> TensorCPU: ...
    def as_tensor(self) -> TensorCPU: ...

    # def at(self, arg0: int) -> numpy.ndarray: ...
    def at(self, arg0: int) -> Any: ...
    def copy_to_external(self, arg0: object) -> None: ...
    def data_ptr(self) -> object: ...
    def is_dense_tensor(self) -> bool: ...
    def layout(self) -> str: ...
    def shape(self) -> List[tuple]: ...
    def __add__(self, *args) -> object: ...
    def __and__(self, *args) -> object: ...
    def __eq__(self, *args) -> object: ...
    def __floordiv__(self, *args) -> object: ...
    def __ge__(self, *args) -> object: ...
    def __getitem__(self, i: int) -> TensorCPU: ...
    def __gt__(self, *args) -> object: ...
    def __le__(self, *args) -> object: ...
    def __len__(self) -> int: ...
    def __lt__(self, *args) -> object: ...
    def __mul__(self, *args) -> object: ...
    def __ne__(self, *args) -> object: ...
    def __neg__(self, *args) -> object: ...
    def __or__(self, *args) -> object: ...
    def __pow__(self, *args) -> object: ...
    def __radd__(self, *args) -> object: ...
    def __rand__(self, *args) -> object: ...
    def __rfloordiv__(self, *args) -> object: ...
    def __rmul__(self, *args) -> object: ...
    def __ror__(self, *args) -> object: ...
    def __rpow__(self, *args) -> object: ...
    def __rsub__(self, *args) -> object: ...
    def __rtruediv__(self, *args) -> object: ...
    def __rxor__(self, *args) -> object: ...
    def __sub__(self, *args) -> object: ...
    def __truediv__(self, *args) -> object: ...
    def __xor__(self, *args) -> object: ...
    @property
    def dtype(self) -> DALIDataType: ...

class TensorListGPU:
    # __hash__: ClassVar[None] = ...

    # @overload
    # def __init__(self, object: capsule, layout: str = ...) -> None: ...
    @overload
    def __init__(self, object, layout: str = ...) -> None: ...
    @overload
    def __init__(self, tl: TensorListGPU, layout: object = ...) -> None: ...
    @overload
    def __init__(self, list_of_tensors: list, layout: str = ...) -> None: ...
    @overload
    def __init__(self, object: object, layout: str = ..., device_id: int = ...) -> None: ...
    @overload
    def __init__(self) -> None: ...
    def as_cpu(self) -> TensorListCPU: ...
    def as_reshaped_tensor(self, arg0: List[int]) -> TensorGPU: ...
    def as_tensor(self) -> TensorGPU: ...
    def at(self, arg0: int) -> TensorGPU: ...
    def copy_to_external(
        self,
        ptr: object,
        cuda_stream: object = ...,
        non_blocking: bool = ...,
        use_copy_kernel: bool = ...,
    ) -> None: ...
    def data_ptr(self) -> object: ...
    def is_dense_tensor(self) -> bool: ...
    def layout(self) -> str: ...
    def shape(self) -> List[tuple]: ...
    def __add__(self, *args) -> object: ...
    def __and__(self, *args) -> object: ...
    def __eq__(self, *args) -> object: ...
    def __floordiv__(self, *args) -> object: ...
    def __ge__(self, *args) -> object: ...
    def __getitem__(self, i: int) -> TensorGPU: ...
    def __gt__(self, *args) -> object: ...
    def __le__(self, *args) -> object: ...
    def __len__(self) -> int: ...
    def __lt__(self, *args) -> object: ...
    def __mul__(self, *args) -> object: ...
    def __ne__(self, *args) -> object: ...
    def __neg__(self, *args) -> object: ...
    def __or__(self, *args) -> object: ...
    def __pow__(self, *args) -> object: ...
    def __radd__(self, *args) -> object: ...
    def __rand__(self, *args) -> object: ...
    def __rfloordiv__(self, *args) -> object: ...
    def __rmul__(self, *args) -> object: ...
    def __ror__(self, *args) -> object: ...
    def __rpow__(self, *args) -> object: ...
    def __rsub__(self, *args) -> object: ...
    def __rtruediv__(self, *args) -> object: ...
    def __rxor__(self, *args) -> object: ...
    def __sub__(self, *args) -> object: ...
    def __truediv__(self, *args) -> object: ...
    def __xor__(self, *args) -> object: ...
    @property
    def dtype(self) -> DALIDataType: ...
